from typing import Union, Dict
import numpy as np

class RandomCropSignal:
    """
    Randomly crop the ECG signal starting from a random point.
    """
    def __init__(self, crop_size=None) -> None:
        """
        :param crop_size: The size of the crop.
        """
        self.crop_size = crop_size

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            data = self._process_dict(data)
        elif isinstance(data, np.ndarray):
            data = self._process_array(data)
        else:
            raise TypeError("Data must be either a numpy array or a dictionary.")
        return data

    def _process_array(self, data: np.ndarray) -> np.ndarray:
        """
        Randomly crop the array to the specified crop size.

        :param data: numpy array of the ECG signal
        :return: Cropped numpy array
        """
        if data.shape[-1] < self.crop_size:
            raise ValueError("Data length is smaller than the crop size.")

        start_idx = np.random.randint(0, data.shape[-1] - self.crop_size + 1)
        return data[..., start_idx:start_idx + self.crop_size]

    def _process_dict(self, data: dict) -> dict:
        """
        Process the dictionary data, cropping based on the crop size and adjusting 'rpeaks'.

        :param data: Dictionary containing 'ecg' and optionally 'rpeaks'
        :return: Cropped dictionary with adjusted annotations
        """
        sample = data["ecg"]
        rpeaks = data["rpeaks"]
        
        sample_length = sample.shape[-1]

        if sample_length < self.crop_size:
            raise ValueError("Data length is smaller than the crop size.")

        start_idx = np.random.randint(0, sample_length - self.crop_size + 1)
        end_idx = start_idx + self.crop_size

        cropped_sample = sample[..., start_idx:end_idx]

        cropped_rpeaks = rpeaks[(rpeaks >= start_idx) & (rpeaks < end_idx)] - start_idx

        data['ecg'] = cropped_sample
        data['rpeaks'] = cropped_rpeaks

        return data
